#!/usr/bin/env python3
"""
RAG (Retrieval-Augmented Generation) integration example.

This example demonstrates how to integrate VectorDB with a language model
to create a RAG system for question-answering with context retrieval.
"""

import asyncio
import httpx
import json
from typing import List, Dict, Any, Optional


class RAGSystem:
    """Simple RAG system using VectorDB and a language model."""

    def __init__(self, vectordb_url: str = "http://localhost:8000", llm_url: str = "http://localhost:11434"):
        self.vectordb_url = vectordb_url
        self.llm_url = llm_url
        self.vectordb_client = httpx.AsyncClient(base_url=vectordb_url, timeout=30.0)
        self.llm_client = httpx.AsyncClient(base_url=llm_url, timeout=30.0)

    async def initialize_knowledge_base(self, documents: List[Dict[str, Any]], table_name: str = "knowledge_base"):
        """Initialize the knowledge base with documents."""
        print("📚 Initializing knowledge base...")

        # Create table
        await self.create_table(table_name)

        # Store documents
        texts = [doc["content"] for doc in documents]
        metadata = [doc.get("metadata", {}) for doc in documents]

        result = await self.store_vectors(table_name, texts, metadata)
        print(f"   Stored {len(documents)} documents in knowledge base")
        return result

    async def create_table(self, table_name: str, vector_dimension: int = 1024):
        """Create a vector table."""
        data = {
            "table_name": table_name,
            "vector_dimension": vector_dimension,
            "overwrite": False
        }
        response = await self.vectordb_client.post("/api/v1/tables", json=data)
        response.raise_for_status()
        return response.json()

    async def store_vectors(self, table_name: str, texts: List[str], metadata: List[Dict[str, Any]]):
        """Store vectors in the database."""
        data = {
            "table_name": table_name,
            "texts": texts,
            "metadata": metadata
        }
        response = await self.vectordb_client.post("/api/v1/store", json=data)
        response.raise_for_status()
        return response.json()

    async def retrieve_context(self, query: str, table_name: str = "knowledge_base",
                             limit: int = 5, threshold: float = 0.0) -> List[Dict[str, Any]]:
        """Retrieve relevant context for a query."""
        data = {
            "query": query,
            "table_name": table_name,
            "limit": limit,
            "threshold": threshold
        }
        response = await self.vectordb_client.post("/api/v1/rag/search", json=data)
        response.raise_for_status()
        return response.json()["context"]

    async def generate_answer(self, query: str, context: List[Dict[str, Any]],
                            model: str = "qwen3:0.6b") -> str:
        """Generate an answer using the language model with context."""
        # Build context string
        context_text = ""
        for i, ctx in enumerate(context, 1):
            context_text += f"Context {i}: {ctx['text']}\n\n"

        # Build prompt
        prompt = f"""Based on the following context, please answer the question.

Context:
{context_text}

Question: {query}

Please provide a comprehensive answer based on the given context. If the context doesn't contain enough information to answer the question, please state that clearly.

Answer:"""

        # Call language model
        llm_data = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.7,
                "max_tokens": 500
            }
        }

        response = await self.llm_client.post("/api/generate", json=llm_data)
        response.raise_for_status()
        result = response.json()
        return result.get("response", "No answer generated")

    async def ask_question(self, question: str, table_name: str = "knowledge_base",
                          context_limit: int = 5, model: str = "qwen3:0.6b") -> Dict[str, Any]:
        """Ask a question and get an answer with context."""
        print(f"❓ Question: {question}")

        # Retrieve relevant context
        print("🔍 Retrieving relevant context...")
        context = await self.retrieve_context(question, table_name, limit=context_limit)

        if not context:
            return {
                "question": question,
                "answer": "I couldn't find relevant information to answer your question.",
                "context": [],
                "context_count": 0
            }

        print(f"   Found {len(context)} relevant documents")

        # Generate answer
        print("🤖 Generating answer...")
        answer = await self.generate_answer(question, context, model)

        return {
            "question": question,
            "answer": answer,
            "context": context,
            "context_count": len(context)
        }

    async def close(self):
        """Close clients."""
        await self.vectordb_client.aclose()
        await self.llm_client.aclose()


async def medical_rag_example():
    """Medical knowledge RAG system example."""
    print("🏥 Medical Knowledge RAG System")
    print("=" * 50)

    rag = RAGSystem()

    # Medical knowledge base
    medical_documents = [
        {
            "content": "Diabetes mellitus is a chronic metabolic disorder characterized by high blood glucose levels. Type 1 diabetes results from autoimmune destruction of pancreatic beta cells, while Type 2 diabetes involves insulin resistance and relative insulin deficiency.",
            "metadata": {
                "category": "endocrinology",
                "condition": "diabetes",
                "type": "overview",
                "source": "medical_textbook"
            }
        },
        {
            "content": "Hypertension, or high blood pressure, is defined as systolic blood pressure ≥140 mmHg or diastolic blood pressure ≥90 mmHg. Risk factors include obesity, sedentary lifestyle, excessive salt intake, and genetic predisposition.",
            "metadata": {
                "category": "cardiology",
                "condition": "hypertension",
                "type": "definition",
                "source": "clinical_guidelines"
            }
        },
        {
            "content": "Asthma is a chronic inflammatory disease of the airways characterized by variable airflow obstruction and bronchial hyperresponsiveness. Common triggers include allergens, exercise, cold air, and respiratory infections.",
            "metadata": {
                "category": "pulmonology",
                "condition": "asthma",
                "type": "overview",
                "source": "medical_textbook"
            }
        },
        {
            "content": "COVID-19 is caused by the SARS-CoV-2 virus. Common symptoms include fever, cough, fatigue, and loss of taste or smell. Complications can include pneumonia, acute respiratory distress syndrome, and multi-organ failure.",
            "metadata": {
                "category": "infectious_disease",
                "condition": "covid19",
                "type": "overview",
                "source": "who_guidelines"
            }
        },
        {
            "content": "Treatment of Type 2 diabetes typically begins with lifestyle modifications including diet and exercise. Metformin is usually the first-line medication. Other options include sulfonylureas, GLP-1 agonists, and SGLT-2 inhibitors.",
            "metadata": {
                "category": "endocrinology",
                "condition": "diabetes",
                "type": "treatment",
                "source": "clinical_guidelines"
            }
        }
    ]

    try:
        # Initialize knowledge base
        print("1. Initializing medical knowledge base...")
        await rag.initialize_knowledge_base(medical_documents, "medical_knowledge")
        print()

        # Test questions
        questions = [
            "What is diabetes and how is it treated?",
            "What are the symptoms of COVID-19?",
            "How is hypertension diagnosed?",
            "What triggers asthma attacks?",
            "What are the risk factors for high blood pressure?"
        ]

        print("2. Testing medical knowledge RAG system...")
        print()

        for question in questions:
            result = await rag.ask_question(question, "medical_knowledge")

            print(f"❓ {result['question']}")
            print(f"💡 {result['answer']}")
            print(f"📚 Used {result['context_count']} context documents")
            print("-" * 50)
            print()

        print("✅ Medical RAG example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await rag.close()


async def customer_support_rag_example():
    """Customer support RAG system example."""
    print("🎧 Customer Support RAG System")
    print("=" * 50)

    rag = RAGSystem()

    # Customer support knowledge base
    support_documents = [
        {
            "content": "To reset your password, click on 'Forgot Password' on the login page. Enter your email address and we'll send you a password reset link. The link expires in 24 hours for security reasons.",
            "metadata": {
                "category": "account",
                "topic": "password_reset",
                "type": "instruction",
                "priority": "high"
            }
        },
        {
            "content": "Our software is compatible with Windows 10/11, macOS 10.15+, and most Linux distributions. Minimum requirements: 8GB RAM, 2GB free disk space, and internet connection for activation.",
            "metadata": {
                "category": "technical",
                "topic": "system_requirements",
                "type": "specification",
                "platforms": ["windows", "macos", "linux"]
            }
        },
        {
            "content": "We offer three subscription plans: Basic ($9.99/month), Professional ($19.99/month), and Enterprise (custom pricing). All plans include core features, with Professional adding advanced analytics and Enterprise offering custom integrations.",
            "metadata": {
                "category": "billing",
                "topic": "subscription_plans",
                "type": "pricing",
                "plans": ["basic", "professional", "enterprise"]
            }
        },
        {
            "content": "Technical support is available 24/7 via email and live chat. Phone support is available Monday-Friday, 9 AM-6 PM EST. Premium support customers get priority response times and dedicated account managers.",
            "metadata": {
                "category": "support",
                "topic": "contact_methods",
                "type": "contact_info",
                "availability": "24/7"
            }
        },
        {
            "content": "To cancel your subscription, go to Account Settings > Subscription > Cancel Subscription. You'll continue to have access until the end of your current billing period. Refunds are processed within 5-7 business days.",
            "metadata": {
                "category": "billing",
                "topic": "cancellation",
                "type": "policy",
                "refund_time": "5-7_days"
            }
        }
    ]

    try:
        # Initialize knowledge base
        print("1. Initializing customer support knowledge base...")
        await rag.initialize_knowledge_base(support_documents, "support_knowledge")
        print()

        # Test customer questions
        customer_questions = [
            "I forgot my password, how can I reset it?",
            "What are the system requirements for your software?",
            "How much does your service cost?",
            "How can I contact technical support?",
            "How do I cancel my subscription?"
        ]

        print("2. Testing customer support RAG system...")
        print()

        for question in customer_questions:
            result = await rag.ask_question(question, "support_knowledge")

            print(f"👤 Customer: {result['question']}")
            print(f"🤖 Support Bot: {result['answer']}")
            print(f"📋 Context used: {result['context_count']} documents")
            print("-" * 50)
            print()

        print("✅ Customer Support RAG example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await rag.close()


async def main():
    """Run all RAG examples."""
    print("VectorDB RAG Integration Examples")
    print("=" * 60)
    print()

    # Run medical RAG example
    await medical_rag_example()
    print()
    print("-" * 60)
    print()

    # Run customer support RAG example
    await customer_support_rag_example()


if __name__ == "__main__":
    asyncio.run(main())